#!/usr/bin/env python3
"""
AFL++ AI 种子生成器 - 主程序入口
"""

import os
import sys
import time
from datetime import datetime

# 添加模块路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from core.analyzer import ProgramAnalyzer
from core.generator import SeedGenerator
from core.saver import SeedSaver
from utils.logger import setup_logging
from config.prompts import build_prompt


class AIAFLSeedGenerator:
    def __init__(self, api_key):
        self.api_key = api_key
        self.analyzer = ProgramAnalyzer()
        self.generator = SeedGenerator(api_key)
        self.saver = SeedSaver()
        self.logger = setup_logging()

        # 统计信息
        self.stats = {
            'total_programs': 0,
            'successful': 0,
            'failed': 0,
            'total_seeds': 0,
            'start_time': None,
            'end_time': None
        }

    def process_single_program(self, program_path):
        """处理单个程序"""
        print(f"\n{'=' * 60}")
        print(f" 分析程序: {program_path}")

        # 分析程序
        analysis = self.analyzer.analyze_program(program_path)

        # 构建程序信息
        if analysis['has_source']:
            with open(program_path, 'r', encoding='utf-8', errors='ignore') as f:
                program_content = f.read()[:5000]  # 限制长度避免token超限
            program_info = f"程序类型: {'C' if analysis['program_type'] == 'c' else 'C++'} 源代码\n```c\n{program_content}\n```"
        else:
            program_info = f"程序类型: 二进制文件\n- 文件类型: {analysis['file_type']}\n- 文件大小: {analysis['file_size']} bytes"
            if analysis['strings']:
                program_info += f"\n- 发现的字符串示例: {', '.join(analysis['strings'][:10])}"

        print(" 生成智能提示词...")
        prompt, expected_count = build_prompt(program_info, analysis, analysis['program_type'])

        # 生成种子
        seeds_text, actual_count = self.generator.generate_seeds(program_path, analysis, program_info)

        if seeds_text:
            # 保存种子
            program_name = os.path.basename(program_path)
            save_result = self.saver.save_seeds(seeds_text, program_path, program_name)

            # 记录响应
            self.logger.log_response(program_path, prompt, seeds_text, save_result['seed_count'])

            print(f" 成功生成 {save_result['seed_count']} 个种子文件到 {save_result['output_dir']}/")

            self.stats['successful'] += 1
            self.stats['total_seeds'] += save_result['seed_count']
            return True
        else:
            print(f"❌ 生成种子失败: {program_path}")
            self.stats['failed'] += 1
            return False

    def process_directory(self, directory_path):
        """处理目录下的所有程序"""
        self.stats['start_time'] = datetime.now()

        if not os.path.exists(directory_path):
            print(f"❌ 目录不存在: {directory_path}")
            return

        # 支持的C/C++扩展名
        supported_extensions = ['.c', '.cpp', '.cc', '.cxx']

        c_programs = []
        for root, dirs, files in os.walk(directory_path):
            for file in files:
                file_ext = os.path.splitext(file)[1].lower()
                if file_ext in supported_extensions:
                    program_path = os.path.join(root, file)
                    c_programs.append(program_path)

        self.stats['total_programs'] = len(c_programs)

        print(f" 找到 {len(c_programs)} 个C/C++程序")

        # 逐个处理程序
        for i, program_path in enumerate(c_programs, 1):
            print(f"\n 处理进度: {i}/{len(c_programs)}")
            success = self.process_single_program(program_path)

            # 添加延迟避免API限制
            if i < len(c_programs):  # 最后一个不需要延迟
                print(" 等待2秒...")
                time.sleep(2)

        self.stats['end_time'] = datetime.now()
        self._print_summary()

    def _print_summary(self):
        """打印汇总报告"""
        duration = self.stats['end_time'] - self.stats['start_time']

        print(f"\n{'=' * 60}")
        print(" 批量处理完成!")
        print(f"{'=' * 60}")
        print(f"   程序总数: {self.stats['total_programs']}")
        print(f"   成功生成: {self.stats['successful']}")
        print(f"   失败数量: {self.stats['failed']}")
        print(f"   种子总数: {self.stats['total_seeds']}")
        print(f"   总耗时: {duration}")
        print(f"   AI响应日志: ai_responses.log")
        print(f"{'=' * 60}")


def main():
    """主函数"""
    API_KEY = 'sk-d92dbab69fb9454c8a047c43142e164f'  # 你的API Key

    if len(sys.argv) != 2:
        print("用法: python main.py <目标目录>")
        print("示例: python main.py test_programs")
        sys.exit(1)

    target_directory = sys.argv[1]

    print(" AFL++ AI 种子批量生成器")
    print(f"目标目录: {target_directory}")

    generator = AIAFLSeedGenerator(API_KEY)

    if os.path.exists(target_directory):
        generator.process_directory(target_directory)
    else:
        print(f"❌ 目标目录不存在: {target_directory}")


if __name__ == "__main__":
    main()